"""流式响应API路由."""

import logging
from typing import Optional, Dict, Any
from uuid import uuid4

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Request, Query
from fastapi.responses import StreamingResponse

from ..auth.jwt_handler import get_current_user_ws, get_current_user
from ..streaming.websocket_manager import websocket_manager
from ..streaming.sse_handler import sse_handler
from ..streaming.stream_processor import stream_processor
from ..callbacks.streaming import create_streaming_callback
from ..chains.generation import ContentGenerationChain

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/streaming", tags=["streaming"])


@router.websocket("/ws/generation/{task_id}")
async def websocket_generation_endpoint(
    websocket: WebSocket,
    task_id: str,
    token: Optional[str] = Query(None)
):
    """WebSocket流式生成端点.
    
    Args:
        websocket: WebSocket连接
        task_id: 任务ID
        token: JWT令牌（可选，通过查询参数传递）
    """
    client_id = None
    
    try:
        # 认证
        user = await get_current_user_ws(websocket, token)
        if not user:
            await websocket.close(code=1008, reason="Authentication failed")
            return
            
        # 建立连接
        client_id = await websocket_manager.connect(websocket, user["user_id"])
        
        # 绑定任务
        await websocket_manager.bind_task(client_id, task_id)
        
        # 主消息循环
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()
            
            # 处理消息
            message_type = data.get("type")
            
            if message_type == "generate":
                # 开始生成
                await handle_generation_request(
                    client_id=client_id,
                    task_id=task_id,
                    data=data,
                    user=user
                )
                
            elif message_type == "cancel":
                # 取消生成
                await handle_cancel_request(client_id, task_id)
                
            elif message_type == "heartbeat":
                # 心跳
                await websocket_manager.handle_message(client_id, data)
                
            else:
                # 其他消息
                await websocket_manager.handle_message(client_id, data)
                
    except WebSocketDisconnect:
        logger.info(f"WebSocket断开: task_id={task_id}, client_id={client_id}")
    except Exception as e:
        logger.error(f"WebSocket错误: task_id={task_id}, error={e}")
    finally:
        if client_id:
            await websocket_manager.disconnect(client_id)


@router.get("/sse/generation/{task_id}")
async def sse_generation_endpoint(
    request: Request,
    task_id: str,
    template_id: str = Query(..., description="模板ID"),
    model: str = Query("gpt-4", description="模型名称"),
    user: dict = Depends(get_current_user)
):
    """SSE流式生成端点.
    
    Args:
        request: FastAPI请求对象
        task_id: 任务ID
        template_id: 模板ID
        model: 模型名称
        user: 当前用户
    """
    try:
        # 创建生成链
        chain = ContentGenerationChain()
        
        # 获取模板变量（从请求体或查询参数）
        variables = {}  # 简化示例
        
        # 创建token生成器
        async def token_generator():
            async for token in chain.generate(template_id, variables, model):
                yield token
                
        # 使用流处理器处理token流
        async def event_generator():
            async for message in stream_processor.process_token_stream(
                task_id=task_id,
                token_generator=token_generator(),
                metadata={
                    "user_id": user["user_id"],
                    "template_id": template_id,
                    "model": model
                }
            ):
                yield message
                
        # 创建SSE响应
        return await sse_handler.create_stream(
            task_id=task_id,
            request=request,
            generator_func=event_generator()
        )
        
    except Exception as e:
        logger.error(f"SSE生成错误: task_id={task_id}, error={e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/cancel/{task_id}")
async def cancel_generation(
    task_id: str,
    user: dict = Depends(get_current_user)
):
    """取消生成任务.
    
    Args:
        task_id: 任务ID
        user: 当前用户
    """
    # 取消WebSocket任务
    ws_cancelled = stream_processor.cancel_task(task_id)
    
    # 取消SSE流
    sse_cancelled = sse_handler.cancel_stream(task_id)
    
    if ws_cancelled or sse_cancelled:
        # 通知相关客户端
        await websocket_manager.send_to_task(task_id, {
            "type": "cancelled",
            "task_id": task_id
        })
        
        return {"status": "success", "message": f"任务 {task_id} 已取消"}
    else:
        raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在或已完成")


@router.get("/status/{task_id}")
async def get_task_status(
    task_id: str,
    user: dict = Depends(get_current_user)
):
    """获取任务状态.
    
    Args:
        task_id: 任务ID
        user: 当前用户
    """
    # 获取WebSocket任务信息
    task_info = stream_processor.get_task_info(task_id)
    
    if task_info:
        return {
            "task_id": task_id,
            "status": task_info.get("status"),
            "token_count": task_info.get("token_count", 0),
            "char_count": task_info.get("char_count", 0),
            "metadata": task_info.get("metadata", {})
        }
        
    # 检查SSE流
    if sse_handler.is_active(task_id):
        return {
            "task_id": task_id,
            "status": "processing",
            "type": "sse"
        }
        
    raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在")


@router.get("/stats")
async def get_streaming_stats(
    user: dict = Depends(get_current_user)
):
    """获取流式服务统计信息.
    
    Args:
        user: 当前用户（需要管理员权限）
    """
    # 检查管理员权限
    if user.get("role") != "admin":
        raise HTTPException(status_code=403, detail="需要管理员权限")
        
    return {
        "websocket": websocket_manager.get_connection_info(),
        "stream_processor": stream_processor.get_statistics(),
        "sse": {
            "active_streams": len(sse_handler.get_active_streams()),
            "stream_ids": sse_handler.get_active_streams()
        }
    }


async def handle_generation_request(
    client_id: str,
    task_id: str,
    data: Dict[str, Any],
    user: Dict[str, Any]
):
    """处理生成请求.
    
    Args:
        client_id: 客户端ID
        task_id: 任务ID
        data: 请求数据
        user: 用户信息
    """
    try:
        # 获取参数
        template_id = data.get("template_id")
        variables = data.get("variables", {})
        model = data.get("model", "gpt-4")
        
        # 创建生成链
        chain = ContentGenerationChain()
        
        # 创建WebSocket回调
        callback = create_streaming_callback(
            websocket_manager=websocket_manager,
            task_id=task_id
        )
        
        # 生成内容
        async def token_generator():
            async for token in chain.generate(
                template_id=template_id,
                variables=variables,
                model_name=model,
                callbacks=[callback]
            ):
                yield token
                
        # 使用流处理器处理
        async for message in stream_processor.process_token_stream(
            task_id=task_id,
            token_generator=token_generator(),
            metadata={
                "user_id": user["user_id"],
                "template_id": template_id,
                "model": model
            }
        ):
            # 发送到任务的所有客户端
            await websocket_manager.send_to_task(task_id, message)
            
    except Exception as e:
        logger.error(f"生成请求处理失败: task_id={task_id}, error={e}")
        await websocket_manager.send_to_task(task_id, {
            "type": "error",
            "error": str(e)
        })


async def handle_cancel_request(client_id: str, task_id: str):
    """处理取消请求.
    
    Args:
        client_id: 客户端ID
        task_id: 任务ID
    """
    # 取消任务
    if stream_processor.cancel_task(task_id):
        await websocket_manager.send_to_task(task_id, {
            "type": "cancelled",
            "task_id": task_id
        })
        logger.info(f"任务已取消: task_id={task_id}, client_id={client_id}")


@router.websocket("/ws/config-sync")
async def websocket_config_sync_endpoint(
    websocket: WebSocket,
    token: Optional[str] = Query(None)
):
    """配置同步WebSocket端点."""
    client_id = None
    
    try:
        # 先建立连接
        await websocket.accept()
        
        # 然后进行认证
        user = await get_current_user_ws(websocket, token)
        if not user:
            await websocket.close(code=1008, reason="Authentication failed")
            return
            
        client_id = f"config_sync_{user['user_id']}_{uuid4().hex[:8]}"
        
        logger.info(f"Config sync WebSocket connected: client_id={client_id}")
        
        # 主消息循环
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()
            
            # 处理消息
            message_type = data.get("type")
            
            if message_type == "auth":
                # 认证消息
                await websocket.send_json({
                    "type": "auth_success",
                    "client_id": client_id
                })
                
            elif message_type == "sync_request":
                # 同步请求
                await websocket.send_json({
                    "type": "sync_response",
                    "data": {
                        "status": "success",
                        "message": "Configuration sync completed"
                    }
                })
                
            elif message_type == "heartbeat":
                # 心跳
                await websocket.send_json({
                    "type": "heartbeat_ack",
                    "timestamp": data.get("timestamp")
                })
                
            else:
                # 其他消息
                logger.info(f"Config sync message: {message_type}")
                
    except WebSocketDisconnect:
        logger.info(f"Config sync WebSocket断开: client_id={client_id}")
    except Exception as e:
        logger.error(f"Config sync WebSocket错误: client_id={client_id}, error={e}")
    finally:
        if client_id:
            logger.info(f"Config sync WebSocket清理: client_id={client_id}")


@router.websocket("/ws/hot-items")
async def websocket_hot_items_endpoint(
    websocket: WebSocket,
    token: Optional[str] = Query(None)
):
    """热点数据WebSocket端点."""
    client_id = None
    
    try:
        # 先建立连接
        await websocket.accept()
        
        # 然后进行认证
        user = await get_current_user_ws(websocket, token)
        if not user:
            await websocket.close(code=1008, reason="Authentication failed")
            return
            
        client_id = f"hot_items_{user['user_id']}_{uuid4().hex[:8]}"
        
        logger.info(f"Hot items WebSocket connected: client_id={client_id}")
        
        # 主消息循环
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()
            
            # 处理消息
            message_type = data.get("type")
            
            if message_type == "subscribe":
                # 订阅热点数据
                await websocket.send_json({
                    "type": "subscribed",
                    "message": "Successfully subscribed to hot items updates"
                })
                
            elif message_type == "heartbeat":
                # 心跳
                await websocket.send_json({
                    "type": "heartbeat_ack",
                    "timestamp": data.get("timestamp")
                })
                
            else:
                # 其他消息
                logger.info(f"Hot items message: {message_type}")
                
    except WebSocketDisconnect:
        logger.info(f"Hot items WebSocket断开: client_id={client_id}")
    except Exception as e:
        logger.error(f"Hot items WebSocket错误: client_id={client_id}, error={e}")
    finally:
        if client_id:
            logger.info(f"Hot items WebSocket清理: client_id={client_id}")